// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <cuda_runtime.h> // @manual
#include <pollux/testing/gtest_utils.h>

#include <pollux/exec/exchange_source.h>
#include <pollux/exec/plan_node_stats.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <tests/exec/utils/HashJoinTestBase.h>
#include <pollux/testing/exec/util/hive_connector_test_base.h>
#include <pollux/testing/exec/util/local_exchange_source.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/experimental/wave/exec/ToWave.h>
#include <pollux/experimental/wave/exec/WaveHiveDataSource.h>
#include <pollux/experimental/wave/exec/tests/utils/FileFormat.h>
#include <pollux/experimental/wave/exec/tests/utils/WaveTestSplitReader.h>
#include <pollux/vector/vector_fuzzer.h>

TURBO_DECLARE_FLAG(int32_t, wave_max_reader_batch_rows);
TURBO_DECLARE_FLAG(int32_t, max_streams_per_driver);
TURBO_DECLARE_FLAG(int32_t, wave_reader_rows_per_tb);

using namespace kumo::pollux;
using namespace kumo::pollux::core;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;

class HashJoinTest : public virtual exec::test::HashJoinTestBase {
 protected:
  void SetUp() override {
    HashJoinTestBase::SetUp();
    if (int device; cudaGetDevice(&device) != cudaSuccess) {
      GTEST_SKIP() << "No CUDA detected, skipping all tests";
    }
    wave::registerWave();
    wave::WaveHiveDataSource::registerConnector();
    wave::test::WaveTestSplitReader::registerTestSplitReader();
    exec::ExchangeSource::factories().clear();
    exec::ExchangeSource::registerFactory(createLocalExchangeSource);
  }

  static void SetUpTestCase() {
    OperatorTestBase::SetUpTestCase();
  }

  void TearDown() override {
    wave::test::Table::dropAll();
    HiveConnectorTestBase::TearDown();
  }
};

TEST_F(HashJoinTest, twoKeys) {
  probeType_ = ROW({"t_k1", "t_k2", "t_data"}, {BIGINT(), BIGINT(), BIGINT()});
  buildType_ = ROW({"u_k1", "u_k2", "u_data"}, {BIGINT(), BIGINT(), BIGINT()});

  auto build = make_row_vector(
      {"u_k1", "u_k2", "u_data"},
      {make_flat_vector<int64_t>(1000, [&](auto r) { return r; }),
       make_flat_vector<int64_t>(1000, [&](auto r) { return r; }),
       make_flat_vector<int64_t>(1000, [&](auto r) { return r; })});
  auto probe = make_row_vector(
      {"t_k1", "t_k2", "t_data"},
      {make_flat_vector<int64_t>(1000, [&](auto r) { return r + 100; }),
       make_flat_vector<int64_t>(1000, [&](auto r) { return r + 100; }),
       make_flat_vector<int64_t>(1000, [&](auto r) { return r + 2; })});

  HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get())
      .numDrivers(1)
      .probeType(probeType_)
      .probeKeys({"t_k1", "t_k2"})
      .probeVectors({probe})

      .buildType(buildType_)
      .buildKeys({"u_k1", "u_k2"})
      .buildVectors({build})
      .injectSpill(false)
      .referenceQuery(
          "SELECT t_k1, t_k2, t_data, u_k1, u_k2, u_data FROM t, u WHERE t_k1 = u_k1 AND t_k2 = u_k2")
      .run();
}

TEST_F(HashJoinTest, manyHits) {
  probeType_ = ROW({"t_k1", "t_k2", "t_data"}, {BIGINT(), BIGINT(), BIGINT()});
  buildType_ = ROW({"u_k1", "u_k2", "u_data"}, {BIGINT(), BIGINT(), BIGINT()});

  int32_t numRepeats = 20;
  auto build = make_row_vector(
      {"u_k1", "u_k2", "u_data"},
      {make_flat_vector<int64_t>(
           15000, [&](auto r) { return (r / numRepeats) * 9; }),
       make_flat_vector<int64_t>(
           15000, [&](auto r) { return (r / numRepeats) * 9; }),
       make_flat_vector<int64_t>(15000, [&](auto r) { return r; })});
  auto probe = make_row_vector(
      {"t_k1", "t_k2", "t_data"},
      {make_flat_vector<int64_t>(1000, [&](auto r) { return r * 3; }),
       make_flat_vector<int64_t>(1000, [&](auto r) { return r * 3; }),
       make_flat_vector<int64_t>(1000, [&](auto r) { return r; })});

  HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get())
      .numDrivers(1)
      .probeType(probeType_)
      .probeKeys({"t_k1", "t_k2"})
      .probeVectors({probe})

      .buildType(buildType_)
      .buildKeys({"u_k1", "u_k2"})
      .buildVectors({build})
      .injectSpill(false)
      .referenceQuery(
          "SELECT t_k1, t_k2, t_data, u_k1, u_k2, u_data FROM t, u WHERE t_k1 = u_k1 AND t_k2 = u_k2")
      .run();
}
